#include "PropertyKDTree.h"
#include <algorithm>
#include <queue>
#include <cmath>

void PropertyKDTree::build(const std::vector<std::pair<Point3D, size_t>>& points) {
    points_ = points;
    std::vector<size_t> indices(points.size());
    for (size_t i = 0; i < points.size(); ++i) {
        indices[i] = i;
    }
    root_ = buildTree(indices, 0);
}

std::shared_ptr<PropertyKDNode> PropertyKDTree::buildTree(
    const std::vector<size_t>& indices,
    int depth) {
    
    if (indices.empty()) return nullptr;
    
    int axis = depth % 3;
    auto n = indices.size();
    
    // 选择中位数作为分割点
    std::vector<size_t> sorted_indices = indices;
    std::nth_element(sorted_indices.begin(), 
                    sorted_indices.begin() + n/2,
                    sorted_indices.end(),
                    [this, axis](size_t i1, size_t i2) {
                        const auto& p1 = points_[i1].first;
                        const auto& p2 = points_[i2].first;
                        return axis == 0 ? p1.x < p2.x :
                               axis == 1 ? p1.y < p2.y :
                                         p1.z < p2.z;
                    });
    
    size_t median_idx = sorted_indices[n/2];
    auto node = std::make_shared<PropertyKDNode>(
        points_[median_idx].first,
        points_[median_idx].second,
        axis);
    
    // 递归构建左右子树
    if (n > 1) {
        std::vector<size_t> left_indices(sorted_indices.begin(), 
                                       sorted_indices.begin() + n/2);
        std::vector<size_t> right_indices(sorted_indices.begin() + n/2 + 1,
                                        sorted_indices.end());
        
        node->left = buildTree(left_indices, depth + 1);
        node->right = buildTree(right_indices, depth + 1);
    }
    
    return node;
}

void PropertyKDTree::findKNearest(
    const Point3D& query,
    int k,
    std::vector<std::pair<double, size_t>>& nearest) const {
    
    nearest.clear();
    if (!root_) return;
    
    findNearestImpl(query, root_, nearest, k);
    
    // 排序结果
    std::sort(nearest.begin(), nearest.end());
}

void PropertyKDTree::findNearestImpl(
    const Point3D& query,
    const std::shared_ptr<PropertyKDNode>& node,
    std::vector<std::pair<double, size_t>>& nearest,
    int k) const {
    
    if (!node) return;
    
    // 计算到当前点的距离
    double dx = query.x - node->point.x;
    double dy = query.y - node->point.y;
    double dz = query.z - node->point.z;
    double dist = dx*dx + dy*dy + dz*dz;
    
    // 更新最近点集
    if (nearest.size() < k) {
        nearest.emplace_back(dist, node->index);
        std::push_heap(nearest.begin(), nearest.end());
    } else if (dist < nearest.front().first) {
        std::pop_heap(nearest.begin(), nearest.end());
        nearest.back() = {dist, node->index};
        std::push_heap(nearest.begin(), nearest.end());
    }
    
    // 确定搜索顺序
    double diff = node->axis == 0 ? dx :
                 node->axis == 1 ? dy : dz;
    
    auto near = diff <= 0 ? node->left : node->right;
    auto far = diff <= 0 ? node->right : node->left;
    
    findNearestImpl(query, near, nearest, k);
    
    if (nearest.size() < k || diff * diff < nearest.front().first) {
        findNearestImpl(query, far, nearest, k);
    }
} 